package com.zhang.二叉树;

import java.util.Comparator;

/**
 * @author 张吉鑫
 * @Description 红黑树
 * @date 2021/3/17 19:41
 */
public class RBTree<E> extends BBST<E> {
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    public RBTree() {
    }

    public RBTree(Comparator comparator) {
        super(comparator);
    }

    /**
     * 添加后调整
     *
     * @param node
     */
    @Override
    protected void afterAdd(Node<E> node) {
        Node<E> parent = node.parent;

        if (parent == null) {
            black(node);
            return;
        }

        /**
         * parent是黑色节点 不用做任何处理
         */
        if (isBlack(parent)) {
            return;
        }

        /**
         * 判断祖父节点是否是黑色
         */
        Node<E> grand = parent.parent;
        if (isBlack(parent.sibling())) {

            grand = red(grand);

            if (parent.isLeftChild()) {//L
                if (node.isLeftChild()) {//LL
                    parent = black(parent);
                } else {//LR
                    node = black(node);
                    rotateLeft(parent);
                }
                rotateRight(grand);
            } else {//R
                if (node.isLeftChild()) {//RL
                    node = black(node);
                    rotateRight(parent);
                } else {//RR
                    parent = black(parent);
                }
                rotateLeft(grand);
            }
        } else {//叔父节点是红色的
            parent = black(parent);
            black(parent.sibling());
            afterAdd(red(grand));
        }


    }

    /**
     * 删除后调整
     *
     * @param node
     */
    @Override
    protected void afterRemove(Node<E> node, Node<E> childNode) {
        //如果删除的节点是红色 直接删除
        if (isRed(node)) {
            return;
        }

        //如果替代的节点是红色的 把替代的节点染成黑色
        if (isRed(childNode)) {
            black(childNode);
            return;
        }

        Node<E> parent = node.parent;
        //如果父节点是空的 说明删除节点就是根节点
        if (parent == null) {
            return;
        }
        //获取删除节点的兄弟节点
        boolean left = parent.left == null || node.isLeftChild();
        Node<E> sibling = left ? parent.right : parent.left;


        //如果兄弟节点在父节点的左边
        if (sibling.isLeftChild()) {
            //如果兄弟节点是红色的
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rotateRight(parent);
                sibling = parent.left;
            }
            //如果兄弟节点两个子节点都是黑色的
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                //如果父类是黑色节点 会产生下溢 再次调用方法 把父类节点当初删除节点处理
                boolean b = colorOf(parent);
                red(sibling);
                black(parent);
                if (b) {
                    afterRemove(parent, null);
                }
            } else {
                //兄弟节点至少有一个红色节点
                boolean b = colorOf(parent);
                if (isRed(sibling.right)){
                    rotateLeft(sibling);
                    sibling = parent.left;
                }
                color(sibling, b);
                black(parent);
                black(sibling.left);
                rotateRight(parent);
            }
        } else {
            //如果兄弟节点是红色的
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rotateLeft(parent);
                sibling = parent.right;
            }
            //如果兄弟节点两个子节点都是黑色的
            if (isBlack(sibling.right) && isBlack(sibling.left)) {
                //如果父类是黑色的 会产生下溢现象 把父类节点当作删除节点 重新调用方法
                boolean b = colorOf(parent);
                red(sibling);
                black(parent);
                if (b) {
                    afterRemove(parent, null);
                }
            } else {
                //兄弟节点至少有一个红色节点
                boolean b = colorOf(parent);
                if (isRed(sibling.left)) {
                    rotateRight(sibling);
                    sibling = parent.right;
                }
                color(sibling, b);
                black(parent);
                black(sibling.right);
                rotateLeft(parent);
            }
        }
    }

    /**
     * 染色的方法
     *
     * @param node
     * @param color
     * @return
     */
    private Node<E> color(Node<E> node, boolean color) {
        if (node == null) {
            return node;
        }
        ((RBNode<E>) node).color = color;
        return node;
    }

    /**
     * 染成红色
     *
     * @param node
     * @return
     */
    private Node<E> red(Node<E> node) {
        return color(node, RED);
    }

    /**
     * 染成黑色
     *
     * @param node
     * @return
     */
    private Node<E> black(Node<E> node) {
        return color(node, BLACK);
    }

    /**
     * 查看节点颜色
     *
     * @param node
     * @return
     */
    private boolean colorOf(Node<E> node) {
        return node == null ? BLACK : ((RBNode<E>) node).color;
    }

    /**
     * 判断节点是否为黑色
     *
     * @param node
     * @return
     */
    private boolean isBlack(Node<E> node) {
        return colorOf(node) == BLACK;
    }

    /**
     * 判断节点是否为红色
     *
     * @param node
     * @return
     */
    private boolean isRed(Node<E> node) {
        return colorOf(node) == RED;
    }

    /**
     * 添加新节点
     *
     * @param element
     * @param parent
     * @return
     */
    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new RBNode<>(element, parent);
    }

    private static class RBNode<E> extends Node<E> {
        boolean color = RED;

        public RBNode(E element, Node<E> parent) {
            super(element, parent);
        }

        @Override
        public String toString() {
            String str = "";
            if (color == RED) {
                str = "R_";
            }
            return str + element.toString();
        }
    }
}
